#include "stdafx.h"
#include "MemoryPool.h"
#include <new>
#include <assert.h>
#include "GlobalDefinition.h"

using namespace std;

#define MEMPOOL_ALIGNMENT 8

// get to the smallest number of n*MEMPOOL_ALIGNMENT which is higher than nSize
// ie, if nSize == 11, then the revised nSize == 16
int Block::Align(int nSize)
{
	return (nSize + (MEMPOOL_ALIGNMENT - 1)) & ~(MEMPOOL_ALIGNMENT - 1);
}

Block::Block(unsigned short _nUnitSize, unsigned short _nInitSize, unsigned short _nGrowSize)
{
	pBlock = NULL; // the initial Memory link is NULL

	// to make the InitSize and GrowSize be the amount of MEMPOOL_ALIGNMENT
	nInitSize = Align(_nInitSize); // set the count of allocated unit when the first time the MPool is set up
	nGrowSize = Align(_nGrowSize); 

	if (_nUnitSize > 4)
	{
		nUnitSize = Align(_nUnitSize); 
	}
	else
	{
		if (_nUnitSize <= 2)
		{
			nUnitSize = 2; // we must place an extra 2 bytes in order to save the "next free unit number"
		}
		else
		{
			nUnitSize = 4; // need it? no?
		}
	}
}

Block::~Block()
{
	MemoryBlock* pMyBlock = pBlock;

	while(pMyBlock)
	{
		MemoryBlock* tempBlock = pMyBlock->pNext;

		// delete the block
		delete pMyBlock;

		pMyBlock = tempBlock;
	}
}

// allocate the memory
void* Block::Alloc()
{
	if (!pBlock)
	{
		// this is the first time to allocate the pool

		// first allocate a memory block with nInitSize units, and initialize it 
		try
		{
			pBlock = new(nInitSize/nUnitSize, nUnitSize) MemoryBlock(nInitSize/nUnitSize, nUnitSize);
		}
		catch(bad_alloc)
		{
			// if alloc fails, return NULL to the caller
			return NULL;
		}
		
		// if the initialize is successful, return the 1st unit to caller.
		// attention that the address of the first unit is &aData

		// other configuration has been done by the constructor of the MemoryBlock
		return (char*)pBlock->aData + (pBlock->nFirst - 1) * nUnitSize;
	}
	
	// here we can assure there at least exist one block
	
	// traverse the whole memory block to see if we can find the free block
	MemoryBlock* pMyBlock = pBlock;
	// the parent for the first one is NULL
	MemoryBlock* pParent = NULL;
	while (pMyBlock && !pMyBlock->nFree)
	{
		pParent = pMyBlock;
		pMyBlock = pMyBlock->pNext;
	}

	if ( pMyBlock )
	{
		// we have found this free block
		// find the free one
		char* pFree = pMyBlock->aData + (pMyBlock->nFirst * nUnitSize);

		// adjust the head to minus the free one
		pMyBlock->nFirst = *((unsigned short*)pFree);

		pMyBlock->nFree--;
		return (void*)pFree;
	}
	else
	{
		// we can not find a free block

		// check whether need to grow
		if (!nGrowSize)
		{
			return NULL;
		}

		// initialize the new block
		try
		{
			pMyBlock = new(nGrowSize / nUnitSize, nUnitSize) MemoryBlock(nGrowSize / nUnitSize, nUnitSize);
		}
		catch (bad_alloc)
		{
			return NULL;
		}

		pParent->pNext = pMyBlock;

		return (void*)(pMyBlock->aData);
	}
}

void Block::Free (void* pFree)
{
	if (pFree == NULL)
	{
		return;
	}

	// we must assure the object to free belongs to the memory pool
	MemoryBlock* pParent = NULL;
	MemoryBlock* pMyBlock = InWhichBlock(pFree, &pParent);

	if (pMyBlock == NULL)
	{
		// the pFree does not belong to this memory pool
		// TODO: throw exceptions
		assert((int)(&"Should never be here!") == 0);
	}

	// find the memory block that pFree is located in
	pMyBlock->nFree++;
	// set the first free unit number to the first 2 bytes of the freed unit
	*((unsigned short*)pFree) = pMyBlock->nFirst;
	// set the new first free unit
	pMyBlock->nFirst = GetUnitIndex(pFree);

	// check whether all the unit in this block is free
	if (pMyBlock->nFree * nUnitSize == pMyBlock->nSize)
	{
		// we must leave the first block existing, although it might not contain anything
		if (pMyBlock != pBlock)
		{
			// all the unit in this block are free, return the block to the process
			// if the block is not the first block, pParent must not be NULL
			pParent->pNext = pMyBlock->pNext;	
			
			DELETE_IF_NOT_NULL(pMyBlock);
		}
	}
}

MemoryBlock* Block::InWhichBlock( void* p, MemoryBlock** pParentBlock )
{
	// the parent for the first block is NULL
	if (pParentBlock)
	{
		*pParentBlock = NULL;
	}

	if (pBlock == NULL)
	{
		// if this block does not exist, just return NULL
		return NULL;
	}

	MemoryBlock* pMyBlock = pBlock;

	// traverse the link to check which memory block the "pFree" is located in 
	while(!( ((unsigned long)(pMyBlock->aData) <= (unsigned long)p) 
		&& ((unsigned long)p <= ((unsigned long)pMyBlock->aData + pMyBlock->nSize))) )
	{
		if (pParentBlock)
		{
			*pParentBlock = pMyBlock;
		}
		
		pMyBlock = pMyBlock->pNext;
	}

	return pMyBlock;
}

int Block::FirstFreeUnitIndex(void* p)
{
	MemoryBlock* pMyBlock = InWhichBlock(p);

	if (!pMyBlock)
	{
		// this should never be happened
		assert(0);
	}

	return pMyBlock->nFirst;
}

int Block::GetFreeUnitCount(void* p)
{
	MemoryBlock* pMyBlock = InWhichBlock(p);

	if (!pMyBlock)
	{
		// this should never be happened
		assert(0);
	}

	return pMyBlock->nFree;
}